#include <iostream>
#include <cmath>
#include <cuComplex.h>

// FFT kernel - 现在处理三组独立数据
__global__ void fft_kernel(cuComplex* data_1, cuComplex* data_2, cuComplex* data_3, int N) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < N) {
        cuComplex temp_1 = make_cuComplex(0.0f, 0.0f);
        cuComplex temp_2 = make_cuComplex(0.0f, 0.0f);
        cuComplex temp_3 = make_cuComplex(0.0f, 0.0f);

        // 执行三组独立的FFT计算
        for (int k = 0; k < N; k++) {
            // 三组独立的旋转因子计算
            float realPart_1 = cosf(-2.0f * M_PI * idx * k / N);
            float realPart_2 = cosf(-2.0f * M_PI * idx * k / N);
            float realPart_3 = cosf(-2.0f * M_PI * idx * k / N);
            
            float imagPart_1 = sinf(-2.0f * M_PI * idx * k / N);
            float imagPart_2 = sinf(-2.0f * M_PI * idx * k / N);
            float imagPart_3 = sinf(-2.0f * M_PI * idx * k / N);
            
            cuComplex twiddle_1 = make_cuComplex(realPart_1, imagPart_1);
            cuComplex twiddle_2 = make_cuComplex(realPart_2, imagPart_2);
            cuComplex twiddle_3 = make_cuComplex(realPart_3, imagPart_3);

            // 三组独立的加权求和
            temp_1 = cuCaddf(temp_1, cuCmulf(data_1[k], twiddle_1));
            temp_2 = cuCaddf(temp_2, cuCmulf(data_2[k], twiddle_2));
            temp_3 = cuCaddf(temp_3, cuCmulf(data_3[k], twiddle_3));
        }

        data_1[idx] = temp_1;
        data_2[idx] = temp_2;
        data_3[idx] = temp_3;
    }
}

// 三模冗余投票函数
cuComplex tmr_vote(cuComplex a, cuComplex b, cuComplex c) {
    // 比较实部
    bool ab_real = (fabs(cuCrealf(a) - cuCrealf(b)) < 1e-12);
    bool ac_real = (fabs(cuCrealf(a) - cuCrealf(c)) < 1e-12);
    bool bc_real = (fabs(cuCrealf(b) - cuCrealf(c)) < 1e-12);
    
    // 比较虚部
    bool ab_imag = (fabs(cuCimagf(a) - cuCimagf(b)) < 1e-12);
    bool ac_imag = (fabs(cuCimagf(a) - cuCimagf(c)) < 1e-12);
    bool bc_imag = (fabs(cuCimagf(b) - cuCimagf(c)) < 1e-12);

    // 实部表决
    float real_result;
    if (ab_real && ac_real) {
        real_result = cuCrealf(a);
    } else if (ab_real) {
        real_result = cuCrealf(a);
    } else if (ac_real) {
        real_result = cuCrealf(a);
    } else if (bc_real) {
        real_result = cuCrealf(b);
    } else {
        // 三个都不一致，取中值
        float vals[3] = {cuCrealf(a), cuCrealf(b), cuCrealf(c)};
        if (vals[0] > vals[1]) std::swap(vals[0], vals[1]);
        if (vals[1] > vals[2]) std::swap(vals[1], vals[2]);
        if (vals[0] > vals[1]) std::swap(vals[0], vals[1]);
        real_result = vals[1];
    }

    // 虚部表决
    float imag_result;
    if (ab_imag && ac_imag) {
        imag_result = cuCimagf(a);
    } else if (ab_imag) {
        imag_result = cuCimagf(a);
    } else if (ac_imag) {
        imag_result = cuCimagf(a);
    } else if (bc_imag) {
        imag_result = cuCimagf(b);
    } else {
        // 三个都不一致，取中值
        float vals[3] = {cuCimagf(a), cuCimagf(b), cuCimagf(c)};
        if (vals[0] > vals[1]) std::swap(vals[0], vals[1]);
        if (vals[1] > vals[2]) std::swap(vals[1], vals[2]);
        if (vals[0] > vals[1]) std::swap(vals[0], vals[1]);
        imag_result = vals[1];
    }

    return make_cuComplex(real_result, imag_result);
}

int main() {
    const int N = 8;
    cuComplex* h_data_1 = new cuComplex[N];
    cuComplex* h_data_2 = new cuComplex[N];
    cuComplex* h_data_3 = new cuComplex[N];
    cuComplex* h_final_result = new cuComplex[N];

    // 初始化三组相同数据
    for (int i = 0; i < N; i++) {
        h_data_1[i] = make_cuComplex(1.0f, 0.0f);
        h_data_2[i] = make_cuComplex(1.0f, 0.0f);
        h_data_3[i] = make_cuComplex(1.0f, 0.0f);
    }

    cuComplex* d_data_1;
    cuComplex* d_data_2;
    cuComplex* d_data_3;
    cudaMalloc((void**)&d_data_1, N * sizeof(cuComplex));
    cudaMalloc((void**)&d_data_2, N * sizeof(cuComplex));
    cudaMalloc((void**)&d_data_3, N * sizeof(cuComplex));
    
    // 拷贝数据到设备
    cudaMemcpy(d_data_1, h_data_1, N * sizeof(cuComplex), cudaMemcpyHostToDevice);
    cudaMemcpy(d_data_2, h_data_2, N * sizeof(cuComplex), cudaMemcpyHostToDevice);
    cudaMemcpy(d_data_3, h_data_3, N * sizeof(cuComplex), cudaMemcpyHostToDevice);

    // 启动三模冗余FFT计算
    fft_kernel<<<1, N>>>(d_data_1, d_data_2, d_data_3, N);

    // 拷贝结果回主机
    cudaMemcpy(h_data_1, d_data_1, N * sizeof(cuComplex), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_data_2, d_data_2, N * sizeof(cuComplex), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_data_3, d_data_3, N * sizeof(cuComplex), cudaMemcpyDeviceToHost);

    // 执行三模冗余表决
    for (int i = 0; i < N; i++) {
        h_final_result[i] = tmr_vote(h_data_1[i], h_data_2[i], h_data_3[i]);
    }

    // 输出最终结果和错误检测信息
    int error_count = 0;
    for (int i = 0; i < N; i++) {
        // 检查三个结果是否一致
        bool real_match = fabs(cuCrealf(h_data_1[i]) - cuCrealf(h_data_2[i])) < 1e-12 &&
                          fabs(cuCrealf(h_data_1[i]) - cuCrealf(h_data_3[i])) < 1e-12;
        bool imag_match = fabs(cuCimagf(h_data_1[i]) - cuCimagf(h_data_2[i])) < 1e-12 &&
                          fabs(cuCimagf(h_data_1[i]) - cuCimagf(h_data_3[i])) < 1e-12;
        
        if (!real_match || !imag_match) {
            error_count++;
            std::cout << "Error detected at index " << i << ": "
                      << "Module1: " << cuCrealf(h_data_1[i]) << "+" << cuCimagf(h_data_1[i]) << "i, "
                      << "Module2: " << cuCrealf(h_data_2[i]) << "+" << cuCimagf(h_data_2[i]) << "i, "
                      << "Module3: " << cuCrealf(h_data_3[i]) << "+" << cuCimagf(h_data_3[i]) << "i, "
                      << "Voted: " << cuCrealf(h_final_result[i]) << "+" << cuCimagf(h_final_result[i]) << "i"
                      << std::endl;
        }
    }

    if (error_count == 0) {
        std::cout << "All three modules produced identical results:" << std::endl;
        for (int i = 0; i < N; i++) {
            std::cout << "Data[" << i << "] = " << cuCrealf(h_final_result[i]) 
                      << " + " << cuCimagf(h_final_result[i]) << "i" << std::endl;
        }
    } else {
        std::cout << "\nFinal results after error correction:" << std::endl;
        for (int i = 0; i < N; i++) {
            std::cout << "Data[" << i << "] = " << cuCrealf(h_final_result[i]) 
                      << " + " << cuCimagf(h_final_result[i]) << "i" << std::endl;
        }
        std::cout << "\nTotal errors detected and corrected: " << error_count << std::endl;
    }

    // 清理资源
    delete[] h_data_1;
    delete[] h_data_2;
    delete[] h_data_3;
    delete[] h_final_result;
    cudaFree(d_data_1);
    cudaFree(d_data_2);
    cudaFree(d_data_3);

    return 0;
}